<script setup lang="ts">
import { computed, h, ref, onBeforeMount, onMounted, nextTick } from "vue";
import { ColumnsType } from "ant-design-vue/es/table";
import { Table, InputSearch, Button, message } from "ant-design-vue";
import { cos_sim, env, pipeline } from "@xenova/transformers";

env.remoteHost = "/models/";
type RawInfo = {
  name: string;
  id: string;
};
const vectors = [] as any[];
const infos = ref([] as string[]);
const pipe = pipeline("feature-extraction", "bge-base-zh-v1.5", {
  progress_callback: (d: any) => {
    console.log(d);
  },
});
const addInfo = (info: string) => {
  infos.value.unshift(info);
};
const buildVector = async () => {
  if (!items.value.length) return;
  const list = items.value;
  loading.value = true;
  vectors.length = list.length;
  addInfo("开始构建向量");
  await nextTick();
  const embedding = await pipe;
  const questions = list.map((item) => item.name);
  const output = (await embedding(questions, {
    pooling: "mean",
    normalize: true,
  })) as any;
  console.log(output);
  questions.forEach((q, i) => {
    vectors[i] = output[i];
  });
  loading.value = false;
  addInfo("向量构建完成");
};
const topK = 10;
const minScore = 0.6;
const loading = ref(false);
const items = ref<RawInfo[]>([]);
onMounted(() => {
  fetch("data.json")
    .then((res) => res.json())
    .then((list) => (items.value = list));
});
const searchKey = ref("");
const columns = [
  {
    title: "序号",
    dataIndex: "index",
    key: "index",
    customRender: (e: { index: number }) => {
      return h("span", {}, e.index + 1);
    },
    width: 84,
  },
  {
    title: "名称",
    dataIndex: "name",
    key: "name",
  }
] as ColumnsType<any>;

const showItems = computed(() => {
  return searchKey.value ? result.value : items.value;
});

const result = ref<RawInfo[]>([]);

const search = async (e: string) => {
  searchKey.value = e || "";
  if (!e) {
    return;
  }
  const embedding = await pipe;
  const [vector] = await embedding([e], {
    pooling: "mean",
    normalize: true,
  });
  if (!vectors.length) {
    await buildVector();
  }
  const scores = vectors.map((q, i) => {
    return {
      score: cos_sim(vector.data, vectors[i].data),
      index: i,
    };
  });
  scores.sort((a, b) => b.score - a.score);
  console.log(scores);
  result.value = scores
    .filter((e) => e.score > minScore)
    .slice(0, topK)
    .map((s) => items.value[s.index]);
  addInfo(
    `搜索到${result.value.length}条记录:topK=${topK} minScore=${minScore}`
  );
};
</script>

<template>
  <div class="container">
    <div class="frow">
      <InputSearch placeholder="请输入搜索内容" @search="search" />
      <Button size="small" @click="buildVector">向量库构建</Button>
    </div>
    <Table :loading="loading" :dataSource="showItems" :columns="columns" />
    <div>
      <div v-for="t in infos" :key="t">
        {{ t }}
      </div>
    </div>
  </div>
</template>

<style scoped>
.container {
  border: none;
  height: 100vh;
  background-color: white;
  width: 600px;
  align-self: center;
  color: black;
}
.frow {
  display: flex;
}
</style>
